#include "handtracker.h"
#include <vector>

using namespace std;
using namespace cv;
using namespace mlgesture;

HandTracker::HandTracker() {}

HandTracker::HandTracker(const Mat &frame, const Rect &hand, const Mat mask, const double tracking_sensity)
{
	init(frame, hand, mask, tracking_sensity);
}

HandTracker::~HandTracker() {}

void HandTracker::init(const Mat &frame, const Rect &hand, const Mat mask, const double tracking_sensity)
								throw(HandGestureException) {
	if (!frame.data)
		throw HandGestureException("Data matrix is empty");

	setTrackingSensity(tracking_sensity);
	lostCount_ = 0;
	trackHistSize_ = 36;
	trackHistRanges_[0] = 0;
	trackHistRanges_[1] = 180;
	trackHistChannel_ = 0;
	trackHistPRanges_ = trackHistRanges_;

	Mat roi = Mat::zeros(frame.rows, frame.cols, CV_8UC1);
	rectangle(roi, hand, Scalar(255), CV_FILLED);

	if (mask.data)
		roi &= mask;
	calcHist(&frame, 1, &trackHistChannel_, roi, trackHist_, 1, &trackHistSize_, &trackHistPRanges_);
	normalize(trackHist_, trackHist_, 0, 255, CV_MINMAX);

	Mat tmp;
	vector<Mat> hsv_channels;
	cvtColor(frame, tmp, CV_BGR2HSV);
	split(tmp, hsv_channels);
	prevValue_ = hsv_channels[2];

	kFilter_.init(4, 4);
	setIdentity(kFilter_.transitionMatrix);
	kFilter_.transitionMatrix.at<float>(0, 2) = 1;
	kFilter_.transitionMatrix.at<float>(1, 3) = 1;

	setIdentity(kFilter_.measurementMatrix);
	setIdentity(kFilter_.processNoiseCov, Scalar::all(1e-3));
	setIdentity(kFilter_.measurementNoiseCov, Scalar::all(1e-2));
	setIdentity(kFilter_.errorCovPost, Scalar::all(1));
	kFilter_.statePost = *(Mat_<float>(4, 1) << hand.x, hand.y, 0, 0);
	handRect_ = hand;
	handLastPos_ = Point(hand.x, hand.y);
	kalmanMeasurement_ = Mat::zeros(4, 1, CV_32F);
}

Rect HandTracker::track(const Mat &frame, const Mat &skin_mask) throw(HandGestureException, HandLostException) {
	if (!trackHist_.data)
		throw HandGestureException("Tracker are not initialized");

	Mat kalman_prediction = kFilter_.predict();
	handRect_.x = (int)kalman_prediction.at<float>(0);
	handRect_.y = (int)kalman_prediction.at<float>(1);

	if (handRect_.x < 0) handRect_.x = 0;
	if (handRect_.y < 0) handRect_.y = 0;
	if (handRect_.x + handRect_.width >= frame.cols) handRect_.x = frame.cols - handRect_.width - 1;
	if (handRect_.y + handRect_.height >= frame.rows) handRect_.y = frame.rows - handRect_.height - 1;

	Mat hsv;
	vector<Mat> hsv_channels;
	Mat frame_tmp = frame.clone();

	cvtColor(frame, hsv, CV_BGR2HSV);
	split(hsv, hsv_channels);
	int zeros_count = 0;
	int points_count = 0;
	int skin_points = 0;

	Size size_tmp = frame_tmp.size();
	for( int i = 0; i < size_tmp.height; i++ ) {
		unsigned char* ptr_frame = frame_tmp.ptr<unsigned char>(i);
		const unsigned char* ptr_mask = skin_mask.ptr<unsigned char>(i);
		const unsigned char* ptr_hsv = hsv_channels[2].ptr<unsigned char>(i);
		const unsigned char* ptr_prev = prevValue_.ptr<unsigned char>(i);
		for( int j = 0; j < size_tmp.width; ++j ) {
			if (!ptr_mask[j] || abs(ptr_hsv[j] - ptr_prev[j]) < 10) {
				ptr_frame[j*3] = 0;
				ptr_frame[j*3+1] = 0;
				ptr_frame[j*3+2] = 0;
			}
		}
	}

	for (int i = handRect_.y; i <= handRect_.y+handRect_.height; i++) {
		for (int j = handRect_.x; j <= handRect_.x+handRect_.width; j++) {
			unsigned char skin = skin_mask.at<unsigned char>(i,j);
			unsigned char value = hsv_channels[2].at<unsigned char>(i,j);
			unsigned char prev = prevValue_.at<unsigned char>(i,j);
			if (skin) {
				skin_points++;
			}
			if (!skin || abs(value - prev) < 10) {
				zeros_count++;
			}
			points_count++;
		}
	}

	if (1.0*zeros_count / points_count < trackingSensity_) {
		calcBackProject(&frame_tmp, 1, &trackHistChannel_, trackHist_, trackBackProject_, &trackHistPRanges_);
		meanShift(trackBackProject_, handRect_, TermCriteria( CV_TERMCRIT_EPS | CV_TERMCRIT_ITER, 4, 2 ));

		kalmanMeasurement_.at<float>(0) = handRect_.x;
		kalmanMeasurement_.at<float>(1) = handRect_.y;
		kalmanMeasurement_.at<float>(2) = handRect_.x - handLastPos_.x;
		kalmanMeasurement_.at<float>(3) = handRect_.y - handLastPos_.y;
		kFilter_.correct(kalmanMeasurement_);
		handLastPos_ = Point(handRect_.x, handRect_.y);
	}
	prevValue_ = hsv_channels[2];

	if (1.0 * skin_points / points_count < 0.2) {
		lostCount_++;
	}
	else {
		lostCount_ = 0;
	}

	if (lostCount_ >= 20) {
		throw HandLostException("Hand lost");
	}

	return handRect_;
}

Rect HandTracker::getHandRect() const throw(HandGestureException) {
	if (!trackHist_.data)
		throw HandGestureException("Tracker are not initialized");
	return handRect_;
}

void HandTracker::setTrackingSensity(double sensity) {
	trackingSensity_ = sensity;
}
